package com.unionpay.upyzt.util;

import org.bouncycastle.asn1.gm.GMNamedCurves;
import org.bouncycastle.asn1.x9.X9ECParameters;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.engines.SM2Engine;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.bouncycastle.crypto.params.ParametersWithRandom;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.math.ec.ECPoint;
import org.bouncycastle.util.encoders.Base64;
import org.bouncycastle.util.encoders.Hex;

import java.security.InvalidKeyException;
import java.security.SecureRandom;
import java.security.Security;

public class Sm2Utils {
    static {
        if (Security.getProvider(BouncyCastleProvider.PROVIDER_NAME) == null) {
            Security.addProvider(new BouncyCastleProvider());
        }
    }

    public static final String SPEC_NAME = "sm2p256v1";
    private static final X9ECParameters x9ECParameters = GMNamedCurves.getByName(SPEC_NAME);
    private static final ECDomainParameters ecDomainParameters = new ECDomainParameters(x9ECParameters.getCurve(), x9ECParameters.getG(), x9ECParameters.getN());

    // 默认 SM2 加密结果拼接模式 C1 || C3 || C2
    private static final SM2Engine.Mode defaultMode = SM2Engine.Mode.C1C3C2;

    /**
     * SM2 公钥加密
     *
     * @param data         待加密字符串
     * @param publicKeyHex SM2 公钥点
     * @param mode         加密结果拼接模式 (C1 || C3 || C2) 或 (C1 || C2 || C3)，默认 (C1 || C3 || C2)
     * @return 密文 byte[]
     * @throws InvalidKeyException invalid public key
     */
    public static byte[] encryptToBytes(String data, String publicKeyHex, SM2Engine.Mode mode) throws InvalidKeyException {
        byte[] pointBytes = Hex.decode(publicKeyHex);
        ECPublicKeyParameters ecPublicKeyParameters = null;
        try {
            ECPoint q = x9ECParameters.getCurve().decodePoint(pointBytes);
            ecPublicKeyParameters = new ECPublicKeyParameters(q, ecDomainParameters);
        } catch (IllegalArgumentException ignored) {
        }

        return encryptToBytes(data, ecPublicKeyParameters, mode);
    }

    /**
     * SM2 公钥加密
     *
     * @param data                  待加密字符串
     * @param ecPublicKeyParameters SM2 公钥参数
     * @param mode                  加密结果拼接模式 (C1 || C3 || C2) 或 (C1 || C2 || C3)，默认 (C1 || C3 || C2)
     * @return 密文 byte[]
     * @throws InvalidKeyException invalid public key
     */
    private static byte[] encryptToBytes(String data, ECPublicKeyParameters ecPublicKeyParameters, SM2Engine.Mode mode)
            throws InvalidKeyException {
        if (StringUtils.isBlank(data)) {
            return null;
        }
        if (ecPublicKeyParameters == null) {
            throw new InvalidKeyException("ECPublicKeyParameters is null");
        }
        SM2Engine sm2Engine = new SM2Engine(mode);
        sm2Engine.init(true, new ParametersWithRandom(ecPublicKeyParameters,
                new SecureRandom()));
        byte[] dataBytes = data.getBytes();
        byte[] cipherBytes;
        try {
            cipherBytes = sm2Engine.processBlock(dataBytes, 0, dataBytes.length);
            return cipherBytes;
        } catch (InvalidCipherTextException e) {
            return null;
        }
    }

    /**
     * SM2 公钥加密
     *
     * @param data         待加密字符串
     * @param publicKeyHex SM2 公钥点
     * @return 密文 byte[]
     * @throws InvalidKeyException invalid public key
     */
    public static byte[] encryptToBytes(String data, String publicKeyHex)
            throws InvalidKeyException {
        return encryptToBytes(data, publicKeyHex, defaultMode);
    }

    /**
     * SM2 公钥加密(返回 BASE64 encode 后的数据)
     *
     * @param data         待加密字符串
     * @param publicKeyHex SM2 公钥点
     * @param mode         加密结果拼接模式 (C1 || C3 || C2) 或 (C1 || C2 || C3)，默认 (C1 || C3 || C2)
     * @return 密文 base64
     * @throws InvalidKeyException invalid public key
     */
    public static String encrypt(String data, String publicKeyHex, SM2Engine.Mode mode)
            throws InvalidKeyException {
        byte[] cipherBytes = encryptToBytes(data, publicKeyHex, mode);
        if (cipherBytes == null) {
            return null;
        }
        return new String(Base64.encode(cipherBytes));
    }

    /**
     * SM2 公钥加密(返回 BASE64 encode 后的数据)
     *
     * @param data         待加密字符串
     * @param publicKeyHex SM2 公钥点
     * @return 密文 base64
     * @throws InvalidKeyException invalid public key
     */
    public static String encrypt(String data, String publicKeyHex)
            throws InvalidKeyException {
        return encrypt(data, publicKeyHex, defaultMode);
    }
}
